# -*- coding: utf-8 -*-
"""
Ollama客户端实现模块

本模块实现了基于Ollama的大语言模型客户端，继承自BaseLLM抽象基类。
Ollama是一个本地运行大语言模型的工具，支持多种开源模型如Llama、Mistral等。

主要特性：
1. 本地部署：模型运行在本地，保护数据隐私
2. 多模型支持：支持多种开源大语言模型
3. 高性能：针对本地硬件优化的推理引擎
4. 易于使用：简单的API接口和配置

使用场景：
- 需要数据隐私保护的应用
- 离线环境下的AI应用
- 成本敏感的项目（无需API费用）
- 需要自定义模型的场景
"""
from typing import Any, Dict, List, Optional
from langchain_community.llms.ollama import Ollama
from loguru import logger
from .base import BaseLLM
from config.model_config import model_config


class OllamaClient(BaseLLM):
    """
    Ollama大语言模型客户端
    
    这个类实现了BaseLLM接口，提供基于Ollama的本地大语言模型服务。
    Ollama允许在本地运行各种开源大语言模型，如Llama2、Mistral、CodeLlama等。
    
    优势：
    1. 数据隐私：所有处理都在本地进行
    2. 成本控制：无需支付API调用费用
    3. 自定义性：可以使用自定义训练的模型
    4. 稳定性：不依赖外部API服务的可用性
    
    Attributes:
        client (Ollama): LangChain的Ollama客户端实例
        base_url (str): Ollama服务的基础URL
    
    Examples:
        >>> # 使用默认配置
        >>> client = OllamaClient()
        >>> response = client.generate("什么是人工智能？")
        
        >>> # 自定义配置
        >>> client = OllamaClient(
        ...     model_name="llama2:7b",
        ...     base_url="http://localhost:11434",
        ...     temperature=0.8
        ... )
    """
    
    def __init__(self, model_name: str = None, base_url: str = None, **kwargs):
        """
        初始化Ollama客户端
        
        Args:
            model_name (str, optional): 模型名称，如"llama2:7b"、"mistral:7b"等
                                      如果未指定，使用配置文件中的默认模型
            base_url (str, optional): Ollama服务地址，默认为"http://localhost:11434"
            **kwargs: 其他配置参数，包括：
                - temperature (float): 生成温度，控制输出的随机性 (0.0-1.0)
                - max_tokens (int): 最大生成token数
                - top_p (float): 核采样参数
                - top_k (int): Top-K采样参数
        
        Raises:
            ConnectionError: 当无法连接到Ollama服务时
            ValueError: 当模型名称无效时
        
        Examples:
            >>> # 使用默认配置
            >>> client = OllamaClient()
            
            >>> # 指定模型和参数
            >>> client = OllamaClient(
            ...     model_name="deepseek-r1:32b",
            ...     temperature=0.7,
            ...     max_tokens=2048
            ... )
        """
        # 使用配置文件中的默认值或用户指定的值
        model_name = model_name or model_config.OLLAMA_MODEL
        base_url = base_url or model_config.OLLAMA_BASE_URL
        
        # 调用父类初始化方法
        super().__init__(model_name, **kwargs)
        
        # 存储base_url以便后续使用
        self.base_url = base_url
        
        try:
            # 创建LangChain的Ollama客户端实例
            self.client = Ollama(
                model=model_name,
                base_url=base_url,
                temperature=kwargs.get('temperature', model_config.TEMPERATURE),
                # 添加其他可能的参数
                num_predict=kwargs.get('max_tokens', model_config.MAX_TOKENS),
                top_p=kwargs.get('top_p', model_config.TOP_P),
                top_k=kwargs.get('top_k', model_config.TOP_K)
            )
            
            logger.info(f"成功初始化Ollama客户端: {model_name} @ {base_url}")
            
        except Exception as e:
            logger.error(f"初始化Ollama客户端失败: {e}")
            raise ConnectionError(f"无法连接到Ollama服务 ({base_url}): {e}")
    
    def generate(self, prompt: str, **kwargs) -> str:
        """
        生成文本响应
        
        使用Ollama模型根据输入提示词生成文本响应。
        
        Args:
            prompt (str): 输入的提示词或问题
            **kwargs: 生成参数，可以覆盖初始化时的设置：
                - temperature (float): 生成温度
                - max_tokens (int): 最大生成token数
                - top_p (float): 核采样参数
                - top_k (int): Top-K采样参数
        
        Returns:
            str: 模型生成的文本响应
        
        Raises:
            ValueError: 当提示词无效时
            RuntimeError: 当模型生成失败时
            ConnectionError: 当Ollama服务不可用时
        
        Examples:
            >>> response = client.generate("解释一下机器学习的基本概念")
            >>> print(response)
            "机器学习是人工智能的一个重要分支..."
            
            >>> # 使用自定义参数
            >>> response = client.generate(
            ...     "写一首关于春天的诗",
            ...     temperature=0.9,
            ...     max_tokens=200
            ... )
        """
        # 验证输入提示词
        if not self.validate_prompt(prompt):
            raise ValueError("提示词无效：不能为空或过长")
        
        try:
            logger.debug(f"开始生成文本，提示词长度: {len(prompt)}")
            
            # 调用Ollama客户端生成文本
            response = self.client.invoke(prompt)
            
            logger.debug(f"文本生成成功，响应长度: {len(response)}")
            return response
            
        except Exception as e:
            error_msg = f"Ollama文本生成失败: {e}"
            logger.error(error_msg)
            raise RuntimeError(error_msg)
    
    def batch_generate(self, prompts: List[str], **kwargs) -> List[str]:
        """
        批量生成文本响应
        
        对多个提示词进行批量处理。由于Ollama目前不支持原生批量API，
        这里采用串行处理的方式，逐个处理每个提示词。
        
        Args:
            prompts (List[str]): 提示词列表
            **kwargs: 生成参数，应用于所有提示词
        
        Returns:
            List[str]: 对应每个提示词的生成结果列表
        
        Raises:
            ValueError: 当提示词列表为空或包含无效元素时
            RuntimeError: 当批量生成过程中出现错误时
        
        Examples:
            >>> prompts = [
            ...     "什么是深度学习？",
            ...     "什么是神经网络？",
            ...     "什么是自然语言处理？"
            ... ]
            >>> responses = client.batch_generate(prompts)
            >>> len(responses) == len(prompts)
            True
        
        Note:
            由于是串行处理，批量生成的时间约等于单次生成时间乘以提示词数量。
            对于大量提示词，建议考虑并发处理或分批处理。
        """
        if not prompts:
            raise ValueError("提示词列表不能为空")
        
        # 验证所有提示词
        for i, prompt in enumerate(prompts):
            if not self.validate_prompt(prompt):
                raise ValueError(f"第{i+1}个提示词无效: {prompt[:50]}...")
        
        logger.info(f"开始批量生成文本，共{len(prompts)}个提示词")
        
        results = []
        failed_count = 0
        
        for i, prompt in enumerate(prompts):
            try:
                logger.debug(f"处理第{i+1}/{len(prompts)}个提示词")
                result = self.generate(prompt, **kwargs)
                results.append(result)
                
            except Exception as e:
                failed_count += 1
                error_msg = f"第{i+1}个提示词生成失败: {e}"
                logger.error(error_msg)
                results.append(f"[生成失败: {e}]")
        
        if failed_count > 0:
            logger.warning(f"批量生成完成，{failed_count}/{len(prompts)}个提示词失败")
        else:
            logger.info(f"批量生成成功完成，共处理{len(prompts)}个提示词")
        
        return results
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        
        返回当前Ollama模型的详细信息，包括模型名称、服务地址、配置参数等。
        
        Returns:
            Dict[str, Any]: 包含模型信息的字典，包含以下字段：
                - name: 模型名称
                - provider: 提供商（"Ollama"）
                - base_url: Ollama服务地址
                - temperature: 生成温度
                - max_tokens: 最大token数
                - top_p: 核采样参数
                - top_k: Top-K采样参数
                - config: 完整配置信息
        
        Examples:
            >>> info = client.get_model_info()
            >>> print(f"当前模型: {info['name']}")
            >>> print(f"服务地址: {info['base_url']}")
        """
        return {
            "name": self.model_name,
            "provider": "Ollama",
            "base_url": self.base_url,
            "temperature": self.config.get('temperature', model_config.TEMPERATURE),
            "max_tokens": self.config.get('max_tokens', model_config.MAX_TOKENS),
            "top_p": self.config.get('top_p', model_config.TOP_P),
            "top_k": self.config.get('top_k', model_config.TOP_K),
            "config": self.get_config(),
            "capabilities": [
                "text_generation",
                "conversation",
                "code_generation",
                "question_answering"
            ],
            "supported_languages": [
                "中文", "English", "日本語", "Français", "Deutsch", "Español"
            ]
        }
    
    def health_check(self) -> bool:
        """
        检查Ollama服务健康状态
        
        通过发送简单的测试请求来验证Ollama服务是否正常运行。
        
        Returns:
            bool: 服务是否健康
        
        Examples:
            >>> if client.health_check():
            ...     print("Ollama服务正常")
            ... else:
            ...     print("Ollama服务异常")
        """
        try:
            # 发送简单的测试请求
            test_response = self.generate("Hello", max_tokens=10)
            return len(test_response.strip()) > 0
            
        except Exception as e:
            logger.warning(f"Ollama健康检查失败: {e}")
            return False